"""
Cross-encoder based reranker to improve retrieval precision.
"""
import os
from typing import List, Dict


class Reranker:
    """Ranks candidate evidence chunks with a cross-encoder model."""

    def __init__(self, enabled: bool = True):
        """
        Initialize reranker.

        Args:
            enabled: Whether reranking is enabled
        """
        self.enabled = enabled and os.getenv("RERANKER_ENABLED", "true").lower() == "true"
        self.model_name = os.getenv(
            "RERANKER_MODEL",
            "cross-encoder/ms-marco-MiniLM-L-6-v2",
        )
        self.model = None

        if self.enabled:
            try:
                from sentence_transformers import CrossEncoder

                self.model = CrossEncoder(self.model_name)
            except ImportError:
                print(
                    "Warning: sentence-transformers not installed. "
                    "Install with: pip install sentence-transformers",
                )
                self.model = None
                self.enabled = False
            except Exception as exc:
                print(f"Failed to load reranker model ({self.model_name}): {exc}")
                self.model = None
                self.enabled = False

    def is_available(self) -> bool:
        """Return True when reranker can be used."""
        return self.enabled and self.model is not None

    def rerank(self, question: str, evidence: List[Dict], top_k: int) -> List[Dict]:
        """
        Rerank evidence based on relevance to the question.

        Args:
            question: User question
            evidence: Candidate evidence list
            top_k: Number of reranked items to return

        Returns:
            Reranked evidence list
        """
        if not self.is_available() or not evidence:
            return evidence[:top_k]

        pairs = [[question, item["chunk"]] for item in evidence]
        scores = self.model.predict(pairs)

        scored = []
        for item, score in zip(evidence, scores):
            reranked_item = dict(item)
            reranked_item["rerank_score"] = float(score)
            scored.append(reranked_item)

        scored.sort(key=lambda x: x["rerank_score"], reverse=True)

        return scored[:top_k]

